#!/usr/bin/env python3
import argparse
import grpc
import os
import sys
from time import sleep

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 'utils/'))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper

def writeIpv4_lpmRules(p4info_helper, sw, match_fields, action_params):
    """Install IPv4 LPM forwarding rule."""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": match_fields
        },
        action_name="MyIngress.ipv4_forward",
        action_params=action_params
    )
    sw.WriteTableEntry(table_entry)


def writeAcl_ipTableRule(p4info_helper, sw, ip_match, priority):
    """Install ACL rule to drop packets by destination IP."""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.acl_table",
        match_fields={
            "hdr.ipv4.dstAddr": ip_match,
        },
        action_name="MyIngress.drop",
        priority=priority
    )
    sw.WriteTableEntry(table_entry)
    print(f"Install ACL rule (by IP) with priority {priority}")

def writeAcl_udpTableRule(p4info_helper, sw, udp_match, priority):
    """Successfully install ACL rule to drop packets by UDP destination port."""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.acl_table",
        match_fields={
            "hdr.udp.dstPort": udp_match
        },
        action_name="MyIngress.drop",
        priority=priority
    )
    sw.WriteTableEntry(table_entry)
    print(f"Successfully install ACL rule (by UDP port) with priority {priority}")


def main(p4info_file_path, bmv2_file_path):
    # Load P4 program info
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Connect to switch s1 via P4Runtime gRPC
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')

        # Establish controller as master
        s1.MasterArbitrationUpdate()
        # Push P4 program to the switch
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)

        # Add L3 forwarding rules
        writeIpv4_lpmRules(p4info_helper, s1, ["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 1})
        writeIpv4_lpmRules(p4info_helper, s1, ["10.0.1.2", 32], {"dstAddr": "08:00:00:00:01:02", "port": 2})
        writeIpv4_lpmRules(p4info_helper, s1, ["10.0.1.3", 32], {"dstAddr": "08:00:00:00:01:03", "port": 3})
        writeIpv4_lpmRules(p4info_helper, s1, ["10.0.1.4", 32], {"dstAddr": "08:00:00:00:01:04", "port": 4})

        # Install ACL rules
        # Drop all packets destined to IP 10.0.1.4
        writeAcl_ipTableRule(p4info_helper, s1,
                             ip_match=("10.0.1.4", 0xFFFFFFFF),  # 32-bit mask = exact match
                             priority=100)

        # Drop all UDP packets to port 80
        writeAcl_udpTableRule(p4info_helper, s1,
                              udp_match=(80, 0xFFFF),  # 16-bit mask = exact match
                              priority=90)
    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)
    ShutdownAllSwitchConnections()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, default='./build/acl.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, default='./build/acl.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print(f"\np4info file not found: {args.p4info}\nHave you run 'make'?")
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print(f"\nBMv2 JSON file not found: {args.bmv2_json}\nHave you run 'make'?")
        parser.exit(1)
    main(args.p4info, args.bmv2_json)
